#include "common.h"
#include <ctype.h>
#include <apr_dbd.h>

/** apr_dbd must be initialized exactly once, so make it a static global
 */
static int dbd_initialized = 0;
static apr_pool_t *apr_dbd_init_pool = NULL;

/** A database connection along with all relevant metadata
 */
typedef struct lua_apr_db_connection {
    const apr_dbd_driver_t *driver;
    apr_dbd_t *handle;
    apr_pool_t *pool;
} lua_apr_db_connection;

/** A database result along with all relevant metadata
 */
typedef struct lua_apr_db_result {
    const apr_dbd_driver_t *driver;
    apr_dbd_results_t *results;
    apr_pool_t *pool;
} lua_apr_db_result;

/** A prepared statement and all metadata
 */
typedef struct lua_apr_db_statement {
    lua_apr_db_connection *conn;
    apr_dbd_prepared_t *stmt;
    apr_pool_t *pool;
    int nargs;
} lua_apr_db_statement;

/** Lua function to connect to a database, returning a particular lua_apr_db_connection userdata
 * \luaparam driver_name the string apr_dbd driver name
 * \luaparam connect_params the parameters for the connection in the driver-specific format for apr_dbd_open
 * \luaparam dbname an optional database to use for the connection
 * \luareturn the lua_apr_db_connection userdata for the connection
 */
static int lua_apr_db_connect(lua_State *L) {
    const char *driver_name = luaL_checkstring(L, 1);
    const char *connect_params = luaL_checkstring(L, 2);
    const char *dbname = lua_tostring(L, 3);
    apr_pool_t *db_pool;
    const apr_dbd_driver_t *db_driver;
    apr_dbd_t *handle;
    apr_status_t rv;
    const char *connect_error;
    int errcode;
    
    rv = apr_pool_create(&db_pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error opening database connection");
    }
    
    rv = apr_dbd_get_driver(db_pool, driver_name, &db_driver);
    if(rv != APR_SUCCESS) {
        apr_pool_destroy(db_pool);
        return luaL_error(L, "Failed to load database driver '%s'", driver_name);
    }
    
    rv = apr_dbd_open(db_driver, db_pool, connect_params, &handle);
    if(rv != APR_SUCCESS) {
        apr_pool_destroy(db_pool);
        return luaL_error(L, "Connection failed to %s database %s", driver_name, connect_params);
    }
    
    if(dbname != NULL) {
        errcode = apr_dbd_set_dbname(db_driver, db_pool, handle, dbname);
        if(errcode) {
            connect_error = apr_dbd_error(db_driver, handle, errcode);
            apr_dbd_close(db_driver, handle);
            apr_pool_destroy(db_pool);
            return luaL_error(L, "Connection failed to %s database %s", driver_name, connect_error);
        }
    }
    
    /*Set up and return the userdata variable to handle the connection lifecycle*/
    lua_apr_db_connection *conn = (lua_apr_db_connection *) lua_newuserdata(L, sizeof(lua_apr_db_connection));
    conn->pool = db_pool;
    conn->driver = db_driver;
    conn->handle = handle;
    
    luaL_getmetatable(L, "apr.db.connection");
    lua_setmetatable(L, -2);
    return 1;
}

/** Disconnect from a database and clear all memory pools - this will be run by the lua GC automatically
 */
static int lua_apr_db_disconnect(lua_State *L) {
    apr_status_t rv;
    
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    if(conn->handle != NULL) {
        rv = apr_dbd_close(conn->driver, conn->handle);
        apr_pool_destroy(conn->pool);
        if(rv != APR_SUCCESS) {
            return luaL_error(L, "Error disconnecting from database");
        }
        
        conn->handle = NULL;
        conn->pool = NULL;
    }
    return 0;
}

/** Free a resultset object and clear all associated memory pools - this will be run by the lua GC
 */
static int lua_apr_db_freeresult(lua_State *L) {
    lua_apr_db_result *res = (lua_apr_db_result *) luaL_checkudata(L, 1, "apr.db.result");
    if(res->results != NULL) {
        apr_pool_destroy(res->pool);
        res->results = NULL;
        res->pool = NULL;
    }
    
    return 0;
}

/** Check that the database connection is still active and valid, returning a boolean
 * \luaparam conn the database connection to check
 * \luareturn a boolean true if the database connection is valid, false otherwise
 */
static int lua_apr_db_check(lua_State *L) {
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    apr_status_t rv;
    
    if(conn->handle == NULL) {
        lua_pushboolean(L, 0);
    } else {
        rv = apr_dbd_check_conn(conn->driver, conn->pool, conn->handle);
        if(rv == APR_SUCCESS) {
            lua_pushboolean(L, 1);
        } else {
            lua_pushboolean(L, 0);
        }
    }
    
    return 1;
}

/** Execute a no-result SQL statement against a connection and return the number of affected rows
 * \luaparam conn the connection to execute the statement against
 * \luaparam sql the full SQL statement to execute
 * \luareturn the number of affected rows
 */
static int lua_apr_db_exec(lua_State *L) {
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    const char *sql = luaL_checkstring(L, 2);
    int nrows = 0;
    int retcode;
    const char *errstr;
    
    retcode = apr_dbd_query(conn->driver, conn->handle, &nrows, sql);
    if(retcode != 0) {
        errstr = apr_dbd_error(conn->driver, conn->handle, retcode);
        return luaL_error(L, "%s executing SQL '%s'", errstr, sql);
    }
    
    lua_pushinteger(L, nrows);
    return 1;
}

/** Execute a result-returning SQL statement against a connection and return the raw userdata for the resultset
 * \luaparam conn the connection to execute the statement against
 * \luaparam sql the full SQL statement to execute
 * \luareturn the lua_apr_db_result struct for the resultset
 */
static int lua_apr_db_select(lua_State *L) {
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    const char *sql = luaL_checkstring(L, 2);
    apr_pool_t *pool;
    apr_status_t rv;
    apr_dbd_results_t *results = NULL;
    int retcode;
    const char *query_error;
    
    if(conn->handle == NULL) {
        return luaL_error(L, "Database connection has already been closed");
    }
    
    rv = apr_pool_create(&pool, conn->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error occurred executing query");
    }
    
    retcode = apr_dbd_select(conn->driver, pool, conn->handle, &results, sql, 1);
    
    if(retcode != 0) {
        apr_pool_destroy(pool);
        query_error = apr_dbd_error(conn->driver, conn->handle, retcode);
        return luaL_error(L, "%s executing SQL '%s'", query_error, sql);
    }
    
    lua_apr_db_result *res = (lua_apr_db_result *) lua_newuserdata(L, sizeof(lua_apr_db_result));
    res->pool = pool;
    res->driver = conn->driver;
    res->results = results;
    
    luaL_getmetatable(L, "apr.db.result");
    lua_setmetatable(L, -2);
    return 1;
}

/** Escape a parameter for inclusion in an SQL statement for a particular database type
 * \luaparam conn the database connection to use for escaping
 * \luaparam in the value to escape
 * \luareturn the escaped value
 */
static int lua_apr_db_escape(lua_State *L) {
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    const char *in = luaL_checkstring(L, 2);
    const char *out;
    apr_pool_t *pool;
    apr_status_t rv;
    
    rv = apr_pool_create(&pool, NULL);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error escaping db value");
    }
    
    out = apr_dbd_escape(conn->driver, pool, in, conn->handle);
    lua_pushstring(L, out);
    
    apr_pool_destroy(pool);
    return 1;
}

/** Return the number of rows existing in a resultset
 * \luaparam res the lua_apr_db_result resultset to check
 * \luareturn the number for rows in res
 */
static int lua_apr_db_numrows(lua_State *L) {
    lua_apr_db_result *res = (lua_apr_db_result *) luaL_checkudata(L, 1, "apr.db.result");
    if(res->results == NULL) {
        return luaL_error(L, "Database result has already been closed");
    }
    
    lua_pushinteger(L, apr_dbd_num_tuples(res->driver, res->results));
    return 1;
}

/** Return the number of columns existing in a resultset
 * \luaparam res the lua_apr_db_result resultset to check
 * \luareturn the number of columns in res
 */
static int lua_apr_db_numcols(lua_State *L) {
    lua_apr_db_result *res = (lua_apr_db_result *) luaL_checkudata(L, 1, "apr.db.result");
    if(res->results == NULL) {
        return luaL_error(L, "Database result has already been closed");
    }
    
    lua_pushinteger(L, apr_dbd_num_cols(res->driver, res->results));
    return 1;
}

/** Closure function that returns the next row from a resultset
 * \luareturn a table for the row data, indexed by either field name or column number depending on the return_hash var
 */
static int lua_apr_db_iterrow(lua_State *L) {
    lua_apr_db_result *res = (lua_apr_db_result *) luaL_checkudata(L, lua_upvalueindex(1), "apr.db.result");
    int return_hash = lua_toboolean(L, lua_upvalueindex(2));
    int numcols = lua_tointeger(L, lua_upvalueindex(3));
    const char *entryval;
    lua_Number entrynum;
    
    apr_pool_t *pool;
    apr_status_t rv;
    apr_dbd_row_t *row = NULL;
    
    rv = apr_pool_create(&pool, res->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error getting database row");
    }

    int retcode = apr_dbd_get_row(res->driver, pool, res->results, &row, -1);
    if(retcode == -1) {
        apr_pool_destroy(pool);
        return 0;
    }

    lua_newtable(L);

    int i;
    for(i=0; i<numcols; i++) {
        if(return_hash) {
            lua_pushstring(L, apr_dbd_get_name(res->driver, res->results, i));
        } else {
            lua_pushinteger(L, i + 1);
        }
        
        entryval = apr_dbd_get_entry(res->driver, row, i);
        /*Try converting to a lua number if possible -- THIS IS A HACK, but apr_dbd does not provide real type info*/
        if(isspace(*entryval) || (*entryval == '0' && *(entryval + 1) != '\0')) {
            /*White-space or zero-padded strings will NOT be considered numbers*/
            lua_pushstring(L, entryval);
        } else {
            entrynum = (lua_Number) atof(entryval);
            if(entrynum != 0 || *entryval == '0') {
                lua_pushnumber(L, entrynum);
            } else {
                lua_pushstring(L, entryval);
            }
        }
        lua_settable(L, -3);
    }

    apr_pool_destroy(pool);
    return 1;
}

/** Iterator function that returns the closure to iterator over rows in a resultset
 * \luareturn the closure function lua_apr_db_iterrow to iterate over each individual row
 */
static int lua_apr_db_iterrows(lua_State *L) {
    lua_apr_db_result *res = (lua_apr_db_result *) luaL_checkudata(L, 1, "apr.db.result");
    if(res->results == NULL) {
        return luaL_error(L, "Database result has already been closed");
    }
    
    lua_pushinteger(L, apr_dbd_num_cols(res->driver, res->results));

    lua_pushcclosure(L, lua_apr_db_iterrow, 3);
    return 1;
}

/** Prepare a query for execution using the apr_dbd prepared query function
 * \luaparam conn the database connection to prepare with
 * \luaparam sql the SQL statement to prepare with apr_dbd placeholders
 * \luaparam expectargs the number of parameters to bind into the query - used for error handling later
 * \luareturn the lua_apr_db_statement struct for the statement
 */
static int lua_apr_db_prepare(lua_State *L) {
    lua_apr_db_connection *conn = (lua_apr_db_connection *) luaL_checkudata(L, 1, "apr.db.connection");
    const char *sql = luaL_checkstring(L, 2);
    int expectargs = luaL_checkint(L, 3);
    int retcode;
    const char *errstr;
    apr_pool_t *pool;
    apr_dbd_prepared_t *stmt = NULL;
    apr_status_t rv;
    
    rv = apr_pool_create(&pool, conn->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Memory error preparing query");
    }
    
    retcode = apr_dbd_prepare(conn->driver, pool, conn->handle, sql, NULL, &stmt);
    if(retcode != 0) {
        errstr = apr_dbd_error(conn->driver, conn->handle, retcode);
        return luaL_error(L, "%s preparing SQL '%s'", errstr, sql);
    }
    
    lua_apr_db_statement *prep = (lua_apr_db_statement *) lua_newuserdata(L, sizeof(lua_apr_db_statement));
    prep->conn = conn;
    prep->pool = pool;
    prep->stmt = stmt;
    prep->nargs = expectargs;
    
    luaL_getmetatable(L, "apr.db.statement");
    lua_setmetatable(L, -2);
    return 1;
}

/** Same as lua_apr_db_exec, but execute a prepared statement with bound parameters, not a raw SQL query
 * \luaparam prep the lua_apr_db_statement userdata to execute
 * \luaparam the bound parameters as a table
 * \luareturn the number of rows affected by executing the statement with the given bound parameters
 */
static int lua_apr_db_pexec(lua_State *L) {
    lua_apr_db_statement *prep = (lua_apr_db_statement *) luaL_checkudata(L, 1, "apr.db.statement");
    luaL_checktype(L, 2, LUA_TTABLE);
    int nargs = lua_objlen(L, 2);
    if(nargs != prep->nargs) {
        return luaL_error(L, "Expecting %d parameters, got %d", prep->nargs, nargs);
    }
    
    int nrows = 0;
    int retcode;
    const char *errstr;
    apr_pool_t *pool;
    apr_status_t rv;
    
    rv = apr_pool_create(&pool, prep->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Memory error executing prepared query");
    }
    
    /*Allocate the args array off of the query pool*/
    char **args = (char **) apr_palloc(pool, (nargs * sizeof(char *)));
    char **iterargs = args;
    lua_pushnil(L);
    while(lua_next(L, 2) != 0) {
        *iterargs = (char *) lua_tostring(L, -1);
        ++iterargs;
        lua_pop(L, 1);
    }
    
    retcode = apr_dbd_pquery(prep->conn->driver, pool, prep->conn->handle, &nrows, prep->stmt, nargs, (const char**) args);
    if(retcode != 0) {
        apr_pool_destroy(pool);
        errstr = apr_dbd_error(prep->conn->driver, prep->conn->handle, retcode);
        return luaL_error(L, "%s executing prepared SQL", errstr);
    }
    
    apr_pool_destroy(pool);
    lua_pushinteger(L, nrows);
    return 1;
}

/** Same as lua_apr_db_select, but query a prepared statement instead of a raw SQL query string
 * \luaparam prep the lua_apr_db_statement struct to execute
 * \luaparam the bound parameters for this execution as a table
 * \luareturn the lua_apr_db_result userdata for the result of the query
 */
static int lua_apr_db_pselect(lua_State *L) {
    lua_apr_db_statement *prep = (lua_apr_db_statement *) luaL_checkudata(L, 1, "apr.db.statement");
    luaL_checktype(L, 2, LUA_TTABLE);
    int nargs = lua_objlen(L, 2);
    if(nargs != prep->nargs) {
        return luaL_error(L, "Expecting %d parameters, got %d", prep->nargs, nargs);
    }
    
    int retcode;
    const char *errstr;
    apr_pool_t *pool;
    apr_pool_t *res_pool;
    apr_status_t rv;
    apr_dbd_results_t *results = NULL;
    
    rv = apr_pool_create(&pool, prep->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Memory error executing prepared query");
    }
    
    rv = apr_pool_create(&res_pool, prep->conn->pool);
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Memory error executing prepared query");
    }
    
    /*Allocate the args array off of the query pool*/
    char **args = (char **) apr_palloc(pool, (nargs * sizeof(char *)));
    char **iterargs = args;
    lua_pushnil(L);
    while(lua_next(L, 2) != 0) {
        *iterargs = (char *) lua_tostring(L, -1);
        ++iterargs;
        lua_pop(L, 1);
    }
    
    retcode = apr_dbd_pselect(prep->conn->driver, res_pool, prep->conn->handle, &results, prep->stmt, 1, nargs, (const char **) args);
    
    if(retcode != 0) {
        apr_pool_destroy(pool);
        errstr = apr_dbd_error(prep->conn->driver, prep->conn->handle, retcode);
        return luaL_error(L, "%s executing prepared SQL", errstr);
    }
    
    lua_apr_db_result *res = (lua_apr_db_result *) lua_newuserdata(L, sizeof(lua_apr_db_result));
    res->pool = res_pool;
    res->driver = prep->conn->driver;
    res->results = results;
    
    luaL_getmetatable(L, "apr.db.result");
    lua_setmetatable(L, -2);

    apr_pool_destroy(pool);
    return 1;
}

static const struct luaL_reg lua_apr_dbd [] = {
    {"connect", lua_apr_db_connect},
    {"disconnect", lua_apr_db_disconnect},
    {"check", lua_apr_db_check},
    {"exec", lua_apr_db_exec},
    {"select", lua_apr_db_select},
    {"escape", lua_apr_db_escape},
    {"freeresult", lua_apr_db_freeresult},
    {"numrows", lua_apr_db_numrows},
    {"numcols", lua_apr_db_numcols},
    {"iterrows", lua_apr_db_iterrows},
    {"prepare", lua_apr_db_prepare},
    {"pexec", lua_apr_db_pexec},
    {"pselect", lua_apr_db_pselect},
    {NULL, NULL}
};

/** Set up the metatables and initialize the apr_dbd libraries
 */
int luaopen_apr_db_raw(lua_State *L) {
    apr_status_t rv;
    luaL_register(L, "apr.db.raw", lua_apr_dbd);
    
    rv = apr_initialize();
    if(rv != APR_SUCCESS) {
        return luaL_error(L, "Error initializing APR");
    }
    
    /*apr_dbd_init must be called exactly once per-process*/
    if(!dbd_initialized++) {
        rv = apr_pool_create(&apr_dbd_init_pool, NULL);
        if(rv != APR_SUCCESS) {
            return luaL_error(L, "Error initializing database libraries");
        }
        
        apr_dbd_init(apr_dbd_init_pool);
        dbd_initialized = 1;
    }
    
    luaL_newmetatable(L, "apr.db.connection");
    lua_pushstring(L, "__gc");
    lua_pushcfunction(L, lua_apr_db_disconnect);
    lua_settable(L, -3);
    lua_pop(L, 1);
    
    luaL_newmetatable(L, "apr.db.result");
    lua_pushstring(L, "__gc");
    lua_pushcfunction(L, lua_apr_db_freeresult);
    lua_settable(L, -3);
    lua_pop(L, 1);
    
    luaL_newmetatable(L, "apr.db.statement");
    lua_pop(L, 1);
    
    return 1;
}
